import sys
import os
import collections
import datetime
import json
import numpy
import pandas
from absl import app, flags
import coin.research.feed_interval_load as fil
import coin.strategy.mm.tool.archive_base as abase

from mpl_finance import candlestick_ohlc
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.backends.backend_pdf import PdfPages
import scipy.stats


def plot_dfdict(df, name, nbars=1000):
  print(f"plot {name}")
  fig = plt.figure(figsize=(10, 8))
  fig.set_facecolor('w')
  gs = gridspec.GridSpec(4, 1, height_ratios=[3, 1, 1, 3])
  axes = []
  axes.append(plt.subplot(gs[0]))
  axes.append(plt.subplot(gs[1], sharex=axes[0]))
  axes.append(plt.subplot(gs[2], sharex=axes[0]))
  axes.append(plt.subplot(gs[3], sharex=axes[0]))
  axes[0].get_xaxis().set_visible(False)
  axes[1].get_xaxis().set_visible(False)
  axes[2].get_xaxis().set_visible(False)

  df['dti'] = df.index
  aggdict = {
      'OPEN_MID': 'first',
      'CLOSE_MID': 'last',
      'HIGH_MID': 'max',
      'LOW_MID': 'min',
      'VOLUME_DOLLAR': 'sum',
      'CLOSE_OI_DOLLAR': 'first',
      'dti': 'first'}
  nsqz = max(1, int(df.shape[0] / nbars))
  df = df.groupby((numpy.arange(df.shape[0])/nsqz).astype(int)).agg(aggdict)
  df.index = df['dti']
  del df['dti']

  ohlc = df[['OPEN_MID', 'HIGH_MID', 'LOW_MID', 'CLOSE_MID']].astype(float).values
  oi = df[['CLOSE_OI_DOLLAR']].astype(float).values

  x = numpy.arange(len(df.index))
  dohlc = numpy.hstack((numpy.reshape(x, (-1, 1)), ohlc))
  candlestick_ohlc(axes[0], dohlc, width=0.5, colorup='r', colordown='b')
  up = df['OPEN_MID'] < df['CLOSE_MID']
  axes[1].bar(x[up], df.loc[up, 'VOLUME_DOLLAR'].to_numpy(), color='r', width=1, align='center')
  axes[1].bar(x[~up], df.loc[~up, 'VOLUME_DOLLAR'].to_numpy(), color='b', width=1, align='center')
  axes[1].set_ylabel('volume($)')
  _xticks = numpy.linspace(0, df.shape[0] - 1, 10).astype(int)
  _xlabels = df.index[_xticks].strftime("%y-%m-%d %Hh")

  seri = (
      df['CLOSE_OI_DOLLAR'].rolling('1D').mean()
      / df['VOLUME_DOLLAR'].rolling('1D').mean()
      / df['VOLUME_DOLLAR'].rolling('1D').count().max()).astype(float).values
  axes[2].plot(seri)
  axes[2].set_xticks(_xticks)
  axes[2].set_xticklabels(_xlabels, rotation=45, minor=False)
  axes[2].set_ylabel('oi/24h volume')

  axes[3].plot(oi)
  axes[3].set_xticks(_xticks)
  axes[3].set_xticklabels(_xlabels, rotation=45, minor=False)
  axes[3].set_ylabel('open_interest($)')
  # plt.tight_layout()
  plt.suptitle(name)
  print(f"plot {name} done")
  return df


def dfmix(dfs, daily_usd_volume=None, volume_multiplier=None):
  means = {}
  dfs2 = []
  for df in dfs:
    rows = []
    prev_row = None
    for _, row in df.iterrows():
      if (prev_row is not None
          and ~numpy.isnan(row['OPEN_MID'])
          and (
              ~numpy.isclose(
                  prev_row['CLOSE_MID'],
                  row['OPEN_MID'], rtol=1e-2, atol=0)
              or ~numpy.isclose(
                  row['HIGH_MID'],
                  row['LOW_MID'], rtol=2e-1, atol=0)
              )
          ):
        for col in ['OPEN_MID', 'HIGH_MID', 'LOW_MID', 'CLOSE_MID']:
          row[col] = prev_row['CLOSE_MID']
        row['VOLUME_DOLLAR'] = 0
        row *= numpy.nan
      elif ~numpy.isnan(row['OPEN_MID']):
        prev_row = row
      rows.append(row)
    df = pandas.DataFrame(rows) # .fillna(method='ffill')
    idx = (
        ~numpy.isclose(
            df['CLOSE_MID'].shift(1),
            df['OPEN_MID'], rtol=1e-2, atol=0)
        & ~numpy.isnan(df['OPEN_MID']))
    df.loc[idx, :] = numpy.nan
    dfs2.append(df)
  dfs = dfs2
  nncnt = numpy.sum([~numpy.isnan(df['OPEN_MID']) for df in dfs], axis=0)
  idx_nn0 = numpy.where(nncnt >= 3)[0][0]
  dfs = [df.iloc[idx_nn0:, :] for df in dfs]
  for col in dfs[0].columns:
    seris = []
    for df in dfs:
      seri = df[col].copy()
      # seri[seri == 0] = numpy.nan
      seris.append(seri)
    means[col] = numpy.nanmean(seris, axis=0)

  df = pandas.DataFrame(means)
  df.index = dfs[0].index
  idxnan = numpy.isnan(df['OPEN_MID'])
  if idxnan.sum() > 0:
    idxnnan0 = numpy.where(~idxnan)[0][0]
    df = df.iloc[idxnnan0:, :]
    idxnan = numpy.isnan(df['OPEN_MID'])
    try:
      assert idxnan.sum() == 0
    except:
      pass
  duration = (df.index.max() - df.index.min())
  days = duration.total_seconds() / 3600 / 24
  if volume_multiplier is not None:
    df['VOLUME_DOLLAR'] *= volume_multiplier
    df['CLOSE_OI_DOLLAR'] *= volume_multiplier
  else:
    assert daily_usd_volume is not None
    day_volume = df['VOLUME_DOLLAR'].sum() / days
    df['VOLUME_DOLLAR'] /= day_volume / float(daily_usd_volume)
    df['CLOSE_OI_DOLLAR'] /= day_volume / float(daily_usd_volume)
  day_volume = df['VOLUME_DOLLAR'].sum() / days
  return df, day_volume


def main(_):
  meas_fmts = [
      ('Futures.Binance.v1', '%s-USDT.PERPETUAL'),
      ('Futures.Okex.v3-swap', '%s-USDT.PERPETUAL'),
      ('Futures.Huobi.v1-linear-swap', '%s-USDT.PERPETUAL'),
      ('Futures.Ftx.v1', '%s-USD.PERPETUAL')
      ]
  for base in flags.FLAGS.base.split(","):
    dfs = {}
    for mea, fmt in meas_fmts:
      symbol = (fmt % base).replace("BCH-", "BCHN-")
      dfdict = fil.load_df_dict(mea, flags.FLAGS.trading_date, symbol=symbol)
      dfdict['CLOSE_OI_DOLLAR'] = (
          dfdict['CLOSE_OPEN_INTEREST'].fillna(method='ffill')
          * (dfdict['VOLUME_DOLLAR'] / dfdict['VOLUME']).fillna(method='ffill'))
      assert dfdict['OPEN_MID'].shape[1] <= 1, symbol
      if len(dfdict['OPEN_MID']) == 0:
        continue
      rowdict = {}
      try:
        for col in ['OPEN_MID', 'HIGH_MID', 'LOW_MID', 'CLOSE_MID', 'VOLUME_DOLLAR', 'CLOSE_OI_DOLLAR']:
          rowdict[col] = dfdict[col].iloc[:,0]
      except:
        continue
      df = pandas.DataFrame(rowdict)
      dfs[mea] = df
    dfaqx, day_volume = dfmix(
        list(dfs.values()),
        daily_usd_volume=flags.FLAGS.daily_usd_volume,
        volume_multiplier=flags.FLAGS.volume_multiplier)
    voldict = {'ccy': base, 'day_volume': day_volume}
    json.dump(voldict, open(f'pic/{base}_AQX.json', 'w'))

    if flags.FLAGS.price_precision is not None:
      for pcol in ['OPEN_MID', 'HIGH_MID', 'LOW_MID', 'CLOSE_MID']:
        dfaqx[pcol] = numpy.round(dfaqx[pcol], flags.FLAGS.price_precision)

    if flags.FLAGS.qty_precision is not None:
      ohlc = dfaqx[['OPEN_MID', 'HIGH_MID', 'LOW_MID', 'CLOSE_MID']].astype(float).values
      dfaqx['VOLUME'] = dfaqx['VOLUME_DOLLAR'] / ohlc.mean(1)
      dfaqx['VOLUME'] = numpy.round(dfaqx['VOLUME'], flags.FLAGS.qty_precision)
      dfaqx['VOLUME_DOLLAR'] = dfaqx['VOLUME'] * ohlc.mean(1)

    with PdfPages(f"pic/{base}.pdf") as pdf:
      plot_dfdict(dfaqx, f"{base}_AQX")
      dfaqx.to_csv(f"pic/{base}_AQX_ohlcv.csv")
      outfile = f'pic/{base}_AQX.png'
      plt.savefig(outfile)
      print(os.path.realpath(outfile))
      pdf.savefig()
      plt.close()
      for mea, df in dfs.items():
        plot_dfdict(df, f"{base}_{mea}")
        # plt.savefig(f'pic/{base}_{mea}.png')
        pdf.savefig()
        plt.close()


# ./pyrunner python/coin/research/aqx_past_ohlc.py --trading_date=20210601-20210611 --base=BTC --daily_usd_volume=1e7
# ./pyrunner python/coin/research/aqx_past_ohlc.py --trading_date=20210601-20210611 --base=BTC --volume_multiplier=1e-2
if __name__ == "__main__":
  flags.DEFINE_string('base', None, '')
  flags.DEFINE_float('daily_usd_volume', None, '')
  flags.DEFINE_float('volume_multiplier', None, '')
  flags.DEFINE_integer('price_precision', None, '')
  flags.DEFINE_integer('qty_precision', None, '')
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  app.run(main)
